package com.raos.example.json;

import dev.langchain4j.data.message.UserMessage;
import dev.langchain4j.model.chat.ChatLanguageModel;
import dev.langchain4j.model.chat.request.ChatRequest;
import dev.langchain4j.model.chat.request.ChatRequestParameters;
import dev.langchain4j.model.chat.request.ResponseFormat;
import dev.langchain4j.model.chat.request.ResponseFormatType;
import dev.langchain4j.model.chat.request.json.JsonObjectSchema;
import dev.langchain4j.model.chat.request.json.JsonSchema;
import dev.langchain4j.model.chat.response.ChatResponse;
import dev.langchain4j.service.AiServices;
import java.util.List;

import jakarta.annotation.Resource;
import lombok.extern.slf4j.Slf4j;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;

/**
 * 结构化输出
 *
 * @author raos
 * @date 2023-03-27
 */
@Slf4j
@RequestMapping("/api/json")
@RestController
public class JsonAPI {

    @Resource(name = "ollamaChatModel")
    ChatLanguageModel chatLanguageModel;

    /**
     * 高层API会话聊天（返回结构化结果）
     */
    @GetMapping("/high/chat")
    public String highChat(@RequestParam(value = "message") String message) {
        log.info("message={}", message);
        PersonService personService = AiServices.create(PersonService.class, chatLanguageModel);
        Person person = personService.extractPerson(message);
        log.info("person={}", person);
        return person.toString();
    }

    /**
     * 底层API会话聊天（返回结构化结果）
     */
    @GetMapping("/low/chat")
    public String lowChat(@RequestParam(value = "message") String message) {
        log.info("message={}", message);
        JsonSchema jsonSchema = JsonSchema.builder()
                .rootElement(JsonObjectSchema.builder()
                        .addIntegerProperty("age", "the person's age")
                        .addIntegerProperty("weight", "the person's weight")
                        .required("age", "weight")
                        .build())
                .build();

        ResponseFormat responseFormat = ResponseFormat.builder()
            .type(ResponseFormatType.JSON)
            .jsonSchema(jsonSchema)
            .build();

        ChatResponse chatResponse = chatLanguageModel.chat(ChatRequest.builder()
            .messages(List.of(UserMessage.from(message)))
            .parameters(ChatRequestParameters.builder().responseFormat(responseFormat).build())
            .build());
        String chat = chatResponse.aiMessage().text();
        log.info("chat=[{}]", chat);
        return chat;
    }
}
